package util;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;

import dataProcess.BugInfoToTextFile;
import ds.BugInfo;
import ds.Constants;

public class BugListProcess {
	
	private ArrayList<ArrayList<BugInfo>> bugSet;
	private ArrayList<BugInfo> allBugs;
	
	private ArrayList<String> labels;
	
	private ArrayList<BugInfo> trainBugs;
	private ArrayList<BugInfo> testBugs;
	
	public BugListProcess(ArrayList<BugInfo> bugs, ArrayList<String> categories){
		this.bugSet = new ArrayList<ArrayList<BugInfo>>();
		this.allBugs = bugs;
		this.labels = categories;
		
		this.trainBugs = new ArrayList<BugInfo>();
		trainBugs.addAll(allBugs);
		this.testBugs = new ArrayList<BugInfo>();
	}
	
	public ArrayList<BugInfo> getTrainBugs(){
		return this.trainBugs;
	}
	
	public ArrayList<BugInfo> getTestBugs(){
		return this.testBugs;
	}
	
	public void run(HashMap<String, String> mergeMap, double testRate){
		mergeCategory(mergeMap);
		splitBugToCatagory();
		divideTrainAndTest(testRate);
		writeTrainAndTestForNTM();
		writeTrainAndTestForLPU();
		System.out.println("divide train and test successfully!");
		for(ArrayList<BugInfo> bugs : bugSet){
			System.out.println(bugs.get(0).label + ":" + bugs.size());
		}
	}
	
	public  void mergeCategory(HashMap<String, String> mergeMap) {	
		for (BugInfo curBug : allBugs) {
			if(mergeMap.containsKey(curBug.label))
				curBug.label = mergeMap.get(curBug.label);
		}
	}
	
	public void splitBugToCatagory(){
		for(int i = 0; i < labels.size(); i++){
			ArrayList<BugInfo> subBugList = new ArrayList<BugInfo>();
			this.bugSet.add(subBugList);
		}
		
		for(BugInfo curBug : allBugs){
			int categoryIndex = this.labels.indexOf(curBug.label);
			bugSet.get(categoryIndex).add(curBug);
		}
	}
	
	public void divideTrainAndTest(double testRate){
		for(ArrayList<BugInfo> curList : this.bugSet){
			if(!curList.get(0).label.equals("")){
				HashSet<Integer> testIndexSet = MathUtil.randomNumProducer(
						(int) (testRate * curList.size()), curList.size());
				Iterator<Integer> tis = testIndexSet.iterator();

				while (tis.hasNext()) {
					this.testBugs.add(curList.get(tis.next()));
				}
			}
		}
		
		this.trainBugs.removeAll(testBugs);
	}
	
	public static ArrayList<ArrayList<BugInfo>> divideDataSet(ArrayList<BugInfo> dataSet, double testRate){
		ArrayList<ArrayList<BugInfo>> result = new ArrayList<ArrayList<BugInfo>>();
		ArrayList<BugInfo> test = new ArrayList<BugInfo>();
		
		HashSet<Integer> testIndexSet = MathUtil.randomNumProducer(
				(int) (testRate * dataSet.size()), dataSet.size());
		Iterator<Integer> tis = testIndexSet.iterator();

		while (tis.hasNext()) {
			test.add(dataSet.get(tis.next()));
		}
		dataSet.removeAll(test);
		
		result.add(dataSet);
		result.add(test);
		return result;
	}
	
	public void writeTrainAndTestForNTM(){
		ArrayList<BugInfo> trainForNTM =new ArrayList<BugInfo>();
		trainForNTM.addAll(trainBugs);
		
		int unlabelIndex = this.labels.indexOf("");
		if(unlabelIndex != -1)
			trainForNTM.removeAll(bugSet.get(unlabelIndex));
		
		BugInfoToTextFile.bugList2FileList(trainForNTM, Constants.trainDirForNTM);
		BugInfoToTextFile.bugList2FileList(testBugs, Constants.testDirForNTM);
	}
	
	public void writeTrainAndTestForLPU(){
		BugInfoToTextFile.bugList2FileList(trainBugs, Constants.trainDirForLPU);
		BugInfoToTextFile.bugList2FileList(testBugs, Constants.testDirForLPU);
	}
	
	public ArrayList<Double> getClassProbablity(){
		ArrayList<Double> pClass = new ArrayList<Double>();
		for(int i = 0; i < this.bugSet.size(); i++){
			double p = ((double)bugSet.get(i).size()) / this.allBugs.size();
			pClass.add(p);
		}
		return pClass;
	}
}
